# -*- coding: utf-8 -*-
import os
import yaml
import re

def replace_first_word_before_space(input_string, new_word):
    # 使用 split() 将字符串分割成两部分
    parts = input_string.split(' ', 1)
    if len(parts) < 2:
        # 如果字符串中没有空格，则无法替换
        return input_string
    else:
        # 替换第一个空格前面的字
        parts[0] = str(new_word)  # 将 new_word 转换为字符串
        # 使用 join() 方法重新组合字符串
        return ' '.join(parts)

def write_and_check_duplicate(filename, text):
    # 打开文件以读取内容
    with open(filename, 'r', encoding='utf-8') as file:
        lines = file.readlines()

    line_number = 0
    total=0
    # 判断是否存在相同的文字并记录行数
    for i, line in enumerate(lines):
        total+=1
        if text in line:
            line_number = i

    # 如果存在相同的文字，则输出文字所在行数
    if line_number!=0:
        print(f"文件中已经存在相同内容，位于第 {line_number} 行。")
    else:
        line_number=total
        # 打开文件以追加内容
        with open(filename, 'a', encoding='utf-8') as file:
            # 写入新的文字并换行
            file.write(text)
        #print("内容已写入文件。")
        #print(f"文件 '{filename}' 总共有 {line_number} 行。")

    return line_number

def change_data(path,key,data):

    with open(path,'r',encoding='utf-8') as file:
        lines = file.readlines()
    lines[3] = data + '\n'
    with open(path, 'w',encoding='utf-8') as file:
        file.writelines(lines)   #将Python中的字典或者列表转化为yaml格式的数据
        file.close()

def convert_labelimg_txt_to_yaml(images_dir, annotations_dir, output_yaml):
    with open(output_yaml, 'w', encoding='utf-8') as yaml_file:
        yaml_file.write('path: ../../'+'\n')
        yaml_file.write('train: '+images_dir+'\n')
        yaml_file.write('val: '+images_dir+'\n')
        yaml_file.write('test: '+images_dir+'\n')
        yaml_file.write('# Classes'+'\n')
        yaml_file.write('names:'+'\n')

    class_mapping={}

    with open('E:\\image4\\classes.txt', 'r', encoding='utf-8') as f:
        i=0
        for line in f.readlines():
            print(i,line)
            class_mapping[line.replace("\n", "")]=i
            i+=1

    print(class_mapping)

    my_dict ={}
    file_dict={}
    class_dict={}
    my_dict2={}

    not_dict=[]

    for filename in os.listdir(path+annotations_dir):
        if filename.endswith(".txt"):
            file_path = os.path.join(path+annotations_dir, filename)

            # 读取标签文件
            with open(file_path, 'r') as file:
                lines = file.readlines()
                file_dict[file_path]=lines

            data_list = lines[0].split()
            name=data_list[0]
            #print('name=',name)
            isclass=False
            for key, value in class_mapping.items():
                    #print('key=',key,' value=',value,' name=',name,str(value)==str(name))
                    if(key!=''):
                        match = (str(value)==str(name))
                        if match:
                            print('key=',key,' value=',value)
                            if(class_dict.get(value)):
                                arr=class_dict.get(value)
                            else:
                                arr=[]
                            arr.append(file_path)
                            arr.append(key)
                            class_dict[value]=arr
                            my_dict2[file_path]=key
                            isclass=True
                            break
            if(not isclass):
                not_dict.append(file_path)

    #print('not_dict===',not_dict)
    '''
    for delfile_path in not_dict:
            try:
                # 删除文件
                os.remove(delfile_path)
                print(f"文件 '{delfile_path}' 已成功删除")
            except OSError as e:
                print(f"删除文件时出错: {e}")
    '''
    sorted_object_names = sorted(class_dict, key=lambda x: int(x))
    #print('class_dict=',class_dict)
    #print('sorted_object_names=',sorted_object_names)

    for key in sorted_object_names:
        arr=class_dict[key]
        #print(arr)
        for key2, value2 in my_dict2.items():
            #print(arr[1],value2)
            if(value2==arr[1]):
                my_dict[key2]=value2


    #print('my_dict=',my_dict)
    #print('file_dict=',file_dict)

    # 替换类别编号
    updated_lines = {}
    i=0
    for key, value in my_dict.items():
        if value not in updated_lines:
            updated_lines[value]=i
            i+=1

    #print('updated_lines=',updated_lines)

    for key0, value0 in my_dict.items():
        name=key0
        #print('name=',name)
        for key, value in file_dict.items():
                    if(key!=''):
                        #print('key=',key)
                        match = re.findall(name.replace("\\", ""),key.replace("\\", ""))
                        #print(name.replace("\\", ""),key.replace("\\", ""),name.replace("\\", "")==key.replace("\\", ""))
                        if match:
                            #print('key=',key,' value=',value,' key0=',key0,' value0=',value0)
                            for key1, value1 in updated_lines.items():
                                match2 = re.findall(key1,value0)
                                if(match2):
                                    print(name.replace("\\", ""),key.replace("\\", ""),name.replace("\\", "")==key.replace("\\", ""))
                                    print('value0=',value0,' key1=',key1,re.findall(key1,key.replace("\\", "")))
                                    with open(key, 'w') as file:
                                        # 将文件内容设置为空字符串
                                        file.write('')
                                    for v in value:
                                        r=replace_first_word_before_space(v,value1)
                                        write_and_check_duplicate(key,r)
                            break



    i=0
    for key1, value1 in updated_lines.items():
            print(i,key1)
            with open(output_yaml, 'a', encoding='utf-8') as yaml_file:
                # 写入新的文字并换行
                yaml_file.write('  '+str(i)+': '+key1+'\n')
                i+=1

    nc=i
    print('nc=======',nc)
    file_path = 'models/yolov5s.yaml'
    line_number = 3  # 假设要修改的行号为 3
    new_line_content = 'nc: '+str(nc)+' # number of classes'

    change_data(file_path, line_number, new_line_content)

path='D:/'
# Example usage:
images_dir = 'cococo/images/train'
annotations_dir = 'cococo/labels/train'
output_yaml = path+'/to/output2.yaml'
convert_labelimg_txt_to_yaml(images_dir, annotations_dir, output_yaml)